#include "primer/trie.h"
#include <stack>
#include <string_view>
#include "common/exception.h"
#include "common/macros.h"

namespace bustub {

template <class T>
auto Trie::Get(std::string_view key) const -> const T * {
  // throw NotImplementedException("Trie::Get is not implemented.");
  // You should walk through the trie to find the node corresponding to the key. If the node doesn't exist, return
  // nullptr. After you find the node, you should use `dynamic_cast` to cast it to `const TrieNodeWithValue<T> *`. If
  // dynamic_cast returns `nullptr`, it means the type of the value is mismatched, and you should return nullptr.
  // Otherwise, return the value.
  size_t key_length = key.size();
  const TrieNode *cur_node = root_.get();
  if (key_length == 0) {
    return dynamic_cast<const TrieNodeWithValue<T> *>(cur_node)->value_.get();
  }

  for (size_t i = 0; i < key_length; ++i) {
    auto t = cur_node->children_.find(key[i]);
    if (t == cur_node->children_.end()) {
      return nullptr;
    }
    cur_node = (*t).second.get();
  }
  if (!cur_node->is_value_node_) {
    return nullptr;
  }

  if (auto derived_node_ptr = dynamic_cast<const TrieNodeWithValue<T> *>(cur_node)) {
    return derived_node_ptr->value_.get();
  }
  return nullptr;
}

template <class T>
auto Trie::Put(std::string_view key, T value) const -> Trie {
  // Note that `T` might be a non-copyable type. Always use `std::move` when creating `shared_ptr` on that value.
  // You should walk through the trie and create new nodes if necessary. If the node corresponding to the key already
  // exists, you should create a new `TrieNodeWithValue`.
  // throw NotImplementedException("Trie::Put is not implemented.");

  // 保存向下搜索过程中的节点

  std::stack<std::shared_ptr<const TrieNode>> node_stack;
  size_t key_length = key.size();
  std::shared_ptr<const TrieNode> cur_node = root_;
  if (root_ == nullptr) {
    std::cout << "i am sb" << std::endl;
  }
  size_t i = 0;
  for (i = 0; i < key_length; ++i) {
    auto t = cur_node->children_.find(key[i]);
    if (t == cur_node->children_.end()) {
      break;
    }
    node_stack.push(cur_node);
    cur_node = (*t).second;
  }

  std::stack<std::shared_ptr<TrieNode>> create_stack;
  std::shared_ptr<T> p_value = std::make_shared<T>(std::move(value));
  if (i == key_length) {
    std::shared_ptr<TrieNode> cnt_node = std::make_shared<TrieNodeWithValue<T>>(cur_node->children_, p_value);
    create_stack.push(cnt_node);
  } else {
    std::shared_ptr<TrieNode> cnt_node = std::shared_ptr(cur_node->Clone());
    create_stack.push(cnt_node);
  }

  for (; i < key_length; ++i) {
    if (i == key_length - 1) {
      std::shared_ptr<TrieNode> new_node = std::make_shared<TrieNodeWithValue<T>>(p_value);
      create_stack.push(new_node);
      ++i;
      break;
    }
    std::shared_ptr<TrieNode> new_node = std::make_shared<TrieNode>();
    create_stack.push(new_node);
  }
  --i;
  std::shared_ptr<TrieNode> last = create_stack.top();
  auto debug_value = (dynamic_cast<TrieNodeWithValue<T> *>(last.get()))->value_;
  create_stack.pop();
  while (!create_stack.empty()) {
    create_stack.top()->children_.insert(std::make_pair(key[i], last));
    i--;
    last = create_stack.top();
    create_stack.pop();
  }
  while (!node_stack.empty()) {
    std::shared_ptr<TrieNode> new_node = std::shared_ptr(node_stack.top()->Clone());
    // new_node->children_.insert(std::make_pair(key[i--],last));
    new_node->children_.find(key[i--])->second = last;
    last = new_node;
    node_stack.pop();
  }
  return Trie(last);
}

auto Trie::Remove(std::string_view key) const -> Trie {
  // throw NotImplementedException("Trie::Remove is not implemented.");

  // You should walk through the trie and remove nodes if necessary. If the node doesn't contain a value any more,
  // you should convert it to `TrieNode`. If a node doesn't have children any more, you should remove it.
  std::stack<std::shared_ptr<const TrieNode>> node_stack;
  size_t i = 0;
  size_t key_length = key.size();
  std::shared_ptr<const TrieNode> p = root_;
  for (; i < key_length; ++i) {
    auto t = p->children_.find(key[i]);
    // 要删除的key不存在，返回原本的trie
    if (t == p->children_.end()) {
      return *this;
    }
    node_stack.push(p);
    p = (*t).second;
  }

  BUSTUB_ASSERT(p->is_value_node_ == true, "Trie Romove Error");
  if (p->children_.empty()) {
    p = nullptr;
  } else {
    p = std::make_shared<TrieNode>(p->children_);
  }

  // 复制node_stack中的每一个节点，然后更改children中的那一个

  std::shared_ptr<TrieNode> t = nullptr;
  while (!node_stack.empty()) {
    t = node_stack.top()->Clone();
    if (p == nullptr) {
      t->children_.erase(key[--i]);
    } else {
      t->children_.find(key[--i])->second = p;
    }
    p = t;
    node_stack.pop();
  }
  return Trie(p);
}

// Below are explicit instantiation of template functions.
//
// Generally people would write the implementation of template classes and functions in the header file. However, we
// separate the implementation into a .cpp file to make things clearer. In order to make the compiler know the
// implementation of the template functions, we need to explicitly instantiate them here, so that they can be picked up
// by the linker.

template auto Trie::Put(std::string_view key, uint32_t value) const -> Trie;
template auto Trie::Get(std::string_view key) const -> const uint32_t *;

template auto Trie::Put(std::string_view key, uint64_t value) const -> Trie;
template auto Trie::Get(std::string_view key) const -> const uint64_t *;

template auto Trie::Put(std::string_view key, std::string value) const -> Trie;
template auto Trie::Get(std::string_view key) const -> const std::string *;

// If your solution cannot compile for non-copy tests, you can remove the below lines to get partial score.

using Integer = std::unique_ptr<uint32_t>;

template auto Trie::Put(std::string_view key, Integer value) const -> Trie;
template auto Trie::Get(std::string_view key) const -> const Integer *;

template auto Trie::Put(std::string_view key, MoveBlocked value) const -> Trie;
template auto Trie::Get(std::string_view key) const -> const MoveBlocked *;

}  // namespace bustub
